import torch.nn as nn
import torch
import timm

# from training.models.mlp_mixer import gmlp_s16_224, mixer_b16_224_miil_in21k
# from training.models.unet import ModifiedUNet

# from timm.models.convnext import convnext_tiny
from training.models.cbam_convnext import convnext_tiny, convnext_small, convnext_large_in22k
# from training.models.convnext import convnext_tiny
import training

__all__ = ['cbam_convnext_train']


class CBAMConvnextTrain(nn.Module):

    def __init__(self, num_classes, drop_rate, pretrained):
        super(CBAMConvnextTrain, self).__init__()

        # self.convnext = convnext_tiny(num_classes=num_classes, drop_rate=drop_rate, pretrained=pretrained)
        self.convnext = convnext_large_in22k(num_classes=num_classes, drop_rate=drop_rate, pretrained=pretrained)


    def forward(self, x):
        x = self.convnext(x)
        return x


def cbam_convnext_train(num_classes=2, drop_rate=0., pretrained=False):
    model = CBAMConvnextTrain(num_classes, drop_rate, pretrained)
    model.default_cfg = model.convnext.default_cfg
    return model

if __name__ == "__main__":

    model = CBAMConvnextTrain(num_classes=2, drop_rate=0., pretrained=False)
    x1 = torch.randn(4, 3, 224, 224)
    # x2 = torch.randn(4, 3, 224, 224)
    regression = model(x1)
    print("regression: ".format(), regression.shape)

    # m = timm.create_model('convnext_tiny', pretrained=False)
    # outfeatures = m.forward_features(x1)  ##直接提取网络分类层之前还未池化的特征
    # print("outfeatures: ".format(), outfeatures.shape)

    # m = timm.create_model('convnext_tiny', features_only=True, output_stride=32, out_indices=(0, 3), pretrained=False)
    # print(f'Feature channels: {m.feature_info.channels()}')
    # print(f'Feature reduction: {m.feature_info.reduction()}')
    # o = m(torch.randn(2, 3, 224, 224))
    # for x in o:
    #     print(x.shape)

    m = timm.create_model('convnext_tiny', features_only=True, pretrained=False)
    print(f'Feature channels: {m.feature_info.channels()}')
    o = m(torch.randn(4, 3, 224, 224))
    for x in o:
        print(x.shape)


